
import mindspore as ms

from mindspore import nn
from mindspore.train import Model
from mindvision.engine.callback import LossMonitor
from model import SPPLeNet5, CaltechDataset, gen_dataset


def train(file_path, num_epochs=80):
    # ms.set_context(device_target='GPU')
    ms.set_context(device_target='CPU')
    dataset_class = CaltechDataset(file_path, train=True)
    dataset = gen_dataset(dataset_class, train='train', batch_size=5)

    # 定义ResNet50网络
    network = SPPLeNet5(num_class=256, num_channel=3, num_layers=5)

    step_size = dataset.get_dataset_size()
    lr = nn.cosine_decay_lr(min_lr=0.00001, max_lr=0.001, total_step=step_size * num_epochs,
                            step_per_epoch=step_size, decay_epoch=num_epochs)

    # 定义优化器和损失函数
    opt = nn.Adam(params=network.trainable_params(), learning_rate=lr)
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    # 加载预训练

    # model = Model(network)
    model = Model(network, loss, opt, metrics={"Accuracy": nn.Accuracy()})

    # 模型训练
    model.train(num_epochs, dataset, callbacks=LossMonitor(lr, 10))
    ms.save_checkpoint(network, "./resNet.ckpt")


if __name__ == '__main__':
    file_path = 'E:/Data/caltech_for_user'
    train(file_path=file_path)
